package com.troila.rpc;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ForwardingClientCall;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MetaDataInterceptor implements ClientInterceptor {

    private static final Metadata.Key<String> METADATA_KEY_TOKEN = Metadata.Key.of("token", Metadata.ASCII_STRING_MARSHALLER);

    private String token;

    public MetaDataInterceptor(String token) {
        this.token = token;
    }

    @Override
    public <Q, P> ClientCall<Q, P> interceptCall(MethodDescriptor<Q, P> methodDescriptor,
                                                 CallOptions callOptions,
                                                 Channel channel) {
        return new ForwardingClientCall.SimpleForwardingClientCall<Q, P>(channel.newCall(methodDescriptor, callOptions)) {
            @Override
            public void start(Listener<P> responseListener, Metadata headers) {
                headers.put(METADATA_KEY_TOKEN, token);
                super.start(responseListener, headers);
            }
        };
    }
}
